# 模型管理器 - 统一管理不同类型的大模型接口
import os
import sys
import logging
import time
from typing import Optional, Dict, List, Any
from config import API_KEY, API_BASE_URL, USE_LOCAL_LLM, LLM_MODEL, ENABLE_FREE_LLM, FREE_LLM_PROVIDER, FREE_LLM_API_KEY

logger = logging.getLogger(__name__)

# 配置Hugging Face国内镜像，避免网络超时问题
try:
    # 配置镜像和缓存路径
    env_vars = {
        'HUGGINGFACE_HUB_CACHE': os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models_cache'),
        'TRANSFORMERS_CACHE': os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models_cache')
    }
    os.environ.update(env_vars)
    
    # 添加镜像URL配置 - 使用国内的HF镜像
    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    logger.info("✅ 已配置Hugging Face国内镜像: https://hf-mirror.com")
except Exception as e:
    logger.warning(f"配置Hugging Face镜像时出错: {e}")

# 尝试导入相关模块
try:
    from openai import OpenAI
except ImportError:
    logger.warning("未安装OpenAI库，将无法使用OpenAI API")
    OpenAI = None

try:
    from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
    import torch
    HAS_LOCAL_LLM_SUPPORT = True
except ImportError:
    logger.warning("未安装transformers库，无法使用本地大模型")
    HAS_LOCAL_LLM_SUPPORT = False

# 导入免费大模型API集成
try:
    from free_llm_api_integration import create_free_llm_interface, create_free_llm_manager
    HAS_FREE_LLM_SUPPORT = True
except ImportError:
    logger.warning("未找到free_llm_api_integration模块")
    HAS_FREE_LLM_SUPPORT = False

class ModelManager:
    """模型管理器，统一管理各种大模型接口"""
    
    def __init__(self):
        """初始化模型管理器"""
        self.openai_client = None
        self.local_llm_pipeline = None
        self.tokenizer = None
        self.model = None
        self.free_llm_manager = None
        self.default_provider = "auto"  # auto, openai, local, free
        
        # 初始化状态
        self.has_openai_api = False
        self.has_local_llm = False
        self.has_free_llm = False
        self.last_used_provider = None
        
        # 初始化各模型接口
        self._init_openai_client()
        if USE_LOCAL_LLM:
            self._init_local_llm()
        if ENABLE_FREE_LLM:
            self._init_free_llm()
        
        # 设置默认提供商优先级
        self._set_default_provider_priority()
        
        logger.info("模型管理器初始化完成")
        logger.info(f"可用的模型提供商: {self.get_available_providers()}")
    
    def _init_openai_client(self):
        """初始化OpenAI客户端"""
        try:
            if OpenAI is not None and API_KEY and API_KEY != "your_api_key_here":
                self.openai_client = OpenAI(
                    api_key=API_KEY,
                    base_url=API_BASE_URL
                )
                self.has_openai_api = True
                logger.info("✅ OpenAI API客户端初始化成功")
            else:
                logger.warning("⚠️ OpenAI API不可用：未配置有效的API密钥或未安装OpenAI库")
        except Exception as e:
            logger.error(f"初始化OpenAI客户端时出错: {str(e)}")
            self.has_openai_api = False
    
    def _init_local_llm(self):
        """初始化本地大模型"""
        try:
            if HAS_LOCAL_LLM_SUPPORT:
                logger.info(f"尝试加载本地模型: {LLM_MODEL}")
                
                # 根据不同的模型类型选择不同的加载方式
                if "llama" in LLM_MODEL.lower() or "mistral" in LLM_MODEL.lower() or "deepseek" in LLM_MODEL.lower():
                     # 这些模型通常可以通过pipeline加载，使用镜像URL
                     self.local_llm_pipeline = pipeline(
                         "text-generation",
                         model=LLM_MODEL,
                         torch_dtype=torch.float16,
                         device_map="auto",
                         max_new_tokens=512,
                         temperature=0.7,
                         model_kwargs={
                             'mirror': 'https://hf-mirror.com',
                             'trust_remote_code': True
                         }
                     )
                else:
                    # 尝试直接加载tokenizer和model
                    # 使用镜像URL加载tokenizer和model
                    self.tokenizer = AutoTokenizer.from_pretrained(
                        LLM_MODEL,
                        mirror='https://hf-mirror.com',
                        trust_remote_code=True
                    )
                    self.model = AutoModelForCausalLM.from_pretrained(
                        LLM_MODEL,
                        torch_dtype=torch.float16,
                        device_map="auto",
                        mirror='https://hf-mirror.com',
                        trust_remote_code=True
                    )
                
                self.has_local_llm = True
                logger.info("✅ 本地大模型加载成功")
            else:
                logger.warning("⚠️ 本地大模型不可用：未安装transformers和torch库")
        except Exception as e:
            logger.error(f"加载本地大模型时出错: {str(e)}")
            self.has_local_llm = False
    
    def _init_free_llm(self):
        """初始化免费大模型API"""
        try:
            if HAS_FREE_LLM_SUPPORT:
                self.free_llm_manager = create_free_llm_manager()
                
                # 如果配置了特定的免费大模型提供商和API密钥，尝试注册
                if FREE_LLM_PROVIDER and FREE_LLM_API_KEY:
                    self.free_llm_manager.register_api(FREE_LLM_PROVIDER, FREE_LLM_API_KEY)
                
                # 检查是否有可用的免费大模型API
                available_apis = self.free_llm_manager.get_available_apis()
                if available_apis:
                    self.has_free_llm = True
                    logger.info(f"✅ 免费大模型API初始化成功，可用提供商: {available_apis}")
                else:
                    logger.warning("⚠️ 未找到可用的免费大模型API")
            else:
                logger.warning("⚠️ 免费大模型API不可用：未找到free_llm_api_integration模块")
        except Exception as e:
            logger.error(f"初始化免费大模型API时出错: {str(e)}")
            self.has_free_llm = False
    
    def _set_default_provider_priority(self):
        """设置默认提供商优先级"""
        # 优先级：OpenAI API > 免费大模型API > 本地大模型
        if self.has_openai_api:
            self.default_provider = "openai"
        elif self.has_free_llm:
            self.default_provider = "free"
        elif self.has_local_llm:
            self.default_provider = "local"
        else:
            self.default_provider = "none"
            logger.warning("⚠️ 没有可用的大模型提供商，将使用模拟响应")
    
    def get_available_providers(self) -> List[str]:
        """获取所有可用的模型提供商"""
        providers = []
        if self.has_openai_api:
            providers.append("openai")
        if self.has_free_llm:
            providers.append("free")
        if self.has_local_llm:
            providers.append("local")
        return providers
    
    def generate(self, prompt: str, provider: str = "auto", conversations: List[Dict] = None, **kwargs) -> str:
        """
        生成模型响应
        
        Args:
            prompt: 用户提示
            provider: 模型提供商，可选值：auto, openai, free, local
            conversations: 对话历史
            **kwargs: 其他参数
            
        Returns:
            生成的响应文本
        """
        # 确定使用的提供商
        if provider == "auto":
            provider = self.default_provider
        
        # 根据提供商选择生成方法
        try:
            if provider == "openai" and self.has_openai_api:
                return self._generate_with_openai(prompt, conversations, **kwargs)
            elif provider == "free" and self.has_free_llm:
                return self._generate_with_free_llm(prompt, conversations, **kwargs)
            elif provider == "local" and self.has_local_llm:
                return self._generate_with_local_llm(prompt, conversations, **kwargs)
            else:
                # 没有可用的模型，返回模拟响应
                logger.warning(f"指定的提供商{provider}不可用，返回模拟响应")
                return self._generate_mock_response(prompt)
        except Exception as e:
            logger.error(f"生成响应时出错: {str(e)}")
            # 出错时返回模拟响应
            return self._generate_mock_response(prompt)
    
    def _generate_with_openai(self, prompt: str, conversations: List[Dict] = None, **kwargs) -> str:
        """使用OpenAI API生成响应"""
        try:
            logger.info("使用OpenAI API生成响应")
            self.last_used_provider = "openai"
            
            # 构建消息列表
            messages = []
            
            # 添加系统提示
            system_prompt = kwargs.get("system_prompt", "你是一个智能助手。")
            messages.append({"role": "system", "content": system_prompt})
            
            # 添加历史对话
            if conversations:
                messages.extend(conversations)
            
            # 添加用户提示
            messages.append({"role": "user", "content": prompt})
            
            # 设置参数
            temperature = kwargs.get("temperature", 0.7)
            max_tokens = kwargs.get("max_tokens", 1024)
            model = kwargs.get("model", "gpt-3.5-turbo")
            
            # 调用API
            response = self.openai_client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens
            )
            
            return response.choices[0].message.content.strip()
        except Exception as e:
            logger.error(f"使用OpenAI API时出错: {str(e)}")
            # 如果OpenAI API失败，尝试回退到其他提供商
            if self.has_free_llm:
                logger.info("OpenAI API失败，尝试使用免费大模型API")
                return self._generate_with_free_llm(prompt, conversations, **kwargs)
            elif self.has_local_llm:
                logger.info("OpenAI API失败，尝试使用本地大模型")
                return self._generate_with_local_llm(prompt, conversations, **kwargs)
            raise
    
    def _generate_with_free_llm(self, prompt: str, conversations: List[Dict] = None, **kwargs) -> str:
        """使用免费大模型API生成响应"""
        try:
            logger.info("使用免费大模型API生成响应")
            self.last_used_provider = "free"
            
            # 获取特定的免费模型提供商
            specific_provider = kwargs.get("free_provider")
            
            # 调用免费大模型API
            response = self.free_llm_manager.generate(
                prompt=prompt,
                api_provider=specific_provider,
                conversations=conversations,
                **kwargs
            )
            
            return response
        except Exception as e:
            logger.error(f"使用免费大模型API时出错: {str(e)}")
            # 如果免费大模型API失败，尝试回退到本地大模型
            if self.has_local_llm:
                logger.info("免费大模型API失败，尝试使用本地大模型")
                return self._generate_with_local_llm(prompt, conversations, **kwargs)
            raise
    
    def _generate_with_local_llm(self, prompt: str, conversations: List[Dict] = None, **kwargs) -> str:
        """使用本地大模型生成响应"""
        try:
            logger.info("使用本地大模型生成响应")
            self.last_used_provider = "local"
            
            # 构建完整的提示
            full_prompt = ""
            
            # 添加系统提示
            system_prompt = kwargs.get("system_prompt", "你是一个智能助手。")
            full_prompt += f"{system_prompt}\n\n"
            
            # 添加历史对话
            if conversations:
                for conv in conversations:
                    role = conv.get("role", "user")
                    content = conv.get("content", "")
                    if role == "system":
                        continue  # 系统提示已经添加
                    full_prompt += f"{role}: {content}\n"
            
            # 添加用户提示
            full_prompt += f"user: {prompt}\nassistant: "
            
            # 设置参数
            max_new_tokens = kwargs.get("max_tokens", 512)
            temperature = kwargs.get("temperature", 0.7)
            
            # 使用pipeline生成
            if self.local_llm_pipeline:
                output = self.local_llm_pipeline(
                    full_prompt,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    do_sample=True
                )
                response = output[0]["generated_text"][len(full_prompt):]
            # 直接使用model生成
            elif self.tokenizer and self.model:
                inputs = self.tokenizer(full_prompt, return_tensors="pt").to(self.model.device)
                output = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    do_sample=True
                )
                response = self.tokenizer.decode(output[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
            else:
                raise Exception("本地大模型未正确初始化")
            
            return response.strip()
        except Exception as e:
            logger.error(f"使用本地大模型时出错: {str(e)}")
            raise
    
    def _generate_mock_response(self, prompt: str) -> str:
        """生成模拟响应"""
        logger.warning("使用模拟响应")
        
        # 简单的模拟响应逻辑
        if '外卖' in prompt:
            return "我理解您想点外卖。请告诉我您想吃哪家餐厅的什么菜品，以及送餐地址。"
        elif '浏览器' in prompt or '网页' in prompt:
            return "我可以帮您操作浏览器。请告诉我您想访问哪个网站，以及需要执行什么操作。"
        elif '工作流' in prompt or '自动化' in prompt:
            return "我可以帮您设计和执行自动化工作流。请详细描述您的工作流目标和可用资源。"
        elif '打开' in prompt and '应用' in prompt:
            app_name = prompt.replace('打开', '').replace('应用', '').strip()
            return f"正在打开应用程序: {app_name}"
        elif '关闭' in prompt and '应用' in prompt:
            app_name = prompt.replace('关闭', '').replace('应用', '').strip()
            return f"正在关闭应用程序: {app_name}"
        elif '天气' in prompt:
            return "我需要知道您想查询哪个城市的天气。"
        elif '时间' in prompt or '日期' in prompt:
            current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
            return f"当前时间是: {current_time}"
        else:
            return "感谢您的提问。由于当前没有可用的大模型，我无法提供更详细的回答。请配置有效的API密钥或安装所需的库以启用完整功能。"
    
    def register_free_llm_api(self, provider: str, api_key: str) -> bool:
        """
        注册新的免费大模型API
        
        Args:
            provider: API提供商
            api_key: API密钥
            
        Returns:
            是否注册成功
        """
        try:
            if self.free_llm_manager:
                success = self.free_llm_manager.register_api(provider, api_key)
                if success:
                    self.has_free_llm = True
                    # 如果这是第一个可用的提供商，更新默认提供商
                    if not self.has_openai_api and not self.has_local_llm:
                        self.default_provider = "free"
                return success
            else:
                logger.error("免费大模型管理器未初始化")
                return False
        except Exception as e:
            logger.error(f"注册免费大模型API时出错: {str(e)}")
            return False
    
    def set_preferred_provider(self, provider: str) -> bool:
        """
        设置首选的模型提供商
        
        Args:
            provider: 模型提供商
            
        Returns:
            是否设置成功
        """
        if provider == "openai" and not self.has_openai_api:
            logger.warning("无法设置OpenAI为首选提供商：OpenAI API不可用")
            return False
        elif provider == "free" and not self.has_free_llm:
            logger.warning("无法设置免费大模型为首选提供商：免费大模型API不可用")
            return False
        elif provider == "local" and not self.has_local_llm:
            logger.warning("无法设置本地大模型为首选提供商：本地大模型不可用")
            return False
        elif provider != "auto" and provider != "openai" and provider != "free" and provider != "local":
            logger.warning(f"未知的模型提供商: {provider}")
            return False
        
        self.default_provider = provider
        logger.info(f"已设置首选模型提供商: {provider}")
        return True

# 全局模型管理器实例
_model_manager_instance = None

def get_model_manager() -> ModelManager:
    """
    获取全局模型管理器实例（单例模式）
    
    Returns:
        ModelManager实例
    """
    global _model_manager_instance
    if _model_manager_instance is None:
        _model_manager_instance = ModelManager()
    return _model_manager_instance

# 创建全局model_manager实例，方便直接导入使用
model_manager = get_model_manager()

# 便捷函数：生成模型响应
def generate_model_response(prompt: str, provider: str = "auto", **kwargs) -> str:
    """
    便捷函数，用于生成模型响应
    
    Args:
        prompt: 用户提示
        provider: 模型提供商
        **kwargs: 其他参数
        
    Returns:
        生成的响应
    """
    manager = get_model_manager()
    return manager.generate(prompt, provider, **kwargs)

# 确保model_manager可以被正确导入和使用
__all__ = ['model_manager', 'get_model_manager', 'generate_model_response', 'ModelManager']

# 为了向后兼容，也可以通过直接导入generate函数来使用
def generate(prompt: str, provider: str = "auto", **kwargs) -> str:
    """
    便捷函数，用于生成模型响应（与generate_model_response相同）
    
    Args:
        prompt: 用户提示
        provider: 模型提供商
        **kwargs: 其他参数
        
    Returns:
        生成的响应
    """
    return generate_model_response(prompt, provider, **kwargs)

# 添加generate到导出列表
__all__.append('generate')